import argparse
import os
from src.utils.get import get_df, get_model
import config.config_ml as config  # 通用配置
import config.config_subgroup as subgroup_config  # 子组配置
from src.utils.save import save_metrics
from src.utils.visualize import plot_confusion_matrix

def main():  
    print("开始评估", flush=True)
    # 解析参数
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, default="rf",
                        choices=["rf", "xgb", "svm", "linear", "gbr"])
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--experiment", type=str, default="classification",
                        choices=["classification", "regression"])
    args = parser.parse_args()

    # ===读取数据===
    df = get_df(args.experiment, args.sensitive)
    print("df.shape:", df.shape)

    # ===选择模型===
    ModelClass = get_model(args.model, args.experiment)
    model = ModelClass(random_state=args.seed)
    model.load(f"results/{args.experiment}/{args.model}/best/best_model_{args.model}.pkl")

    # === 输出目录 ===
    output_dir = f"results_subgroup/{args.experiment}/{args.model}"
    os.makedirs(output_dir, exist_ok=True)

    # === 自动亚组分析 ===
    for subgroup_name, subgroup_values in subgroup_config.SUBGROUPS.items():
        subgroup_dir = os.path.join(output_dir, subgroup_name)
        os.makedirs(subgroup_dir, exist_ok=True)
        print(f"\n=== 分析亚组类型: {subgroup_name} ===")

        for subgroup_value in subgroup_values:
            test_df = df[df[subgroup_name] == subgroup_value]
            if len(test_df) == 0:
                print(f"Subgroup {subgroup_value} 测试集为空，跳过...")
                continue

            metrics = model.evaluate(test_df)
            # 移除训练时间和最佳参数信息
            for key in ["train_time", "best_params"]:
                metrics.pop(key, None)

            metrics["subgroup"] = subgroup_value
            metrics["n_samples"] = len(test_df)

            # 保存 CSV
            filtered_metrics = {k: v for k, v in metrics.items() 
                                if k not in ["confusion_matrix", "classification_report"]}
            save_metrics(filtered_metrics, os.path.join(subgroup_dir, f"metrics_{args.model}_{subgroup_value}.csv"))

            # 可视化（仅分类任务）
            if args.experiment == "classification":
                plot_confusion_matrix(metrics['confusion_matrix'], os.path.join(subgroup_dir, f"confusion_matrix_{subgroup_value}.png"))
                print(metrics['confusion_matrix'])
                print(metrics['classification_report'])

if __name__ == "__main__":
    main()
